/*=============================================================================
*
*   Copyright (C) 2020 All rights reserved.
*
*   Filename:		implicit_solver.hpp
*
*   Author: Wang Zhecheng - wangzhecheng@yeah.net
*
*   Date: 2020-03-05 16:24
*
*   Last Editors: Wang Zhecheng - wangzhecheng@yeah.net
*
*   Last modified:	2020-08-04 01:25
*
*   Description:
*
=============================================================================*/
#ifndef IMPLICIT_SOLVER_H
#define IMPLICIT_SOLVER_H

#include "geometry/shape/nodal_dof.hpp"
#include "model/solve/assemble_equation.hpp"
#include "model/solve/solve_time_step.hpp"
#include "model/solve/update_time_step.hpp"
#include "model/time_step.hpp"

namespace geoxfem {
namespace model {

struct implicit_solver {};

template <>
struct TimeStep<implicit_solver> /*{{{*/
    : util::Indexed, ITimePeriod {
  using TimeFunc = arithmetic::Time_Function_Interface;

  TimeStep(const int& nStep, const _real& st, const _real& during)
      : util::Indexed{nStep},
        ITimePeriod{st, during},
        m_tol{0},
        m_converged{true},
        m_kay_mat{},
        m_rhs{},
        m_answer{} {}
  TimeStep()                = default;
  TimeStep(const TimeStep&) = default;
  TimeStep(TimeStep&&)      = default;
  ~TimeStep()               = default;
  TimeStep& operator=(const TimeStep&) = default;
  TimeStep& operator=(TimeStep&&) = default;
  //////////////////////////////////////////////////////////////
  //////////////////////////////////////////////////////////////

  bool ifConverged() const {
    return m_converged;
  }
  void setConverged(const bool& value) {
    m_converged = value;
  }

  _real getTolerance() const {
    return m_tol;
  }

  Eigen::SparseMatrix<_real>& getSparseKay() {
    return m_kay_mat;
  }
  VectorXR& getRhs() {
    return m_rhs;
  }
  VectorXR& getAnswer() {
    return m_answer;
  }

  void solve() {
    GEOXFEM_LOG_SEV(info) << "Solving equation...";
    Eigen::BiCGSTAB<SpMat, Eigen::IncompleteLUT<_real>> solver(m_kay_mat);
    /* Eigen::SimplicialLDLT<SpMat> solver(m_kay_mat); */
    m_answer = solver.solve(m_rhs);
    GEOXFEM_ASSERT(!m_answer.hasNaN());
  }

private:
  _real                      m_tol;
  bool                       m_converged;
  Eigen::SparseMatrix<_real> m_kay_mat;
  VectorXR                   m_rhs;
  VectorXR                   m_answer;

  template <typename CouplePolicy, typename SolvePolicy>
  friend struct update_iteration;
}; /*}}}*/

template <typename MethodPolicy, typename CouplePolicy, typename LoadPolicy>
struct solve_time_step<MethodPolicy, /*{{{*/
                       CouplePolicy,
                       LoadPolicy,
                       implicit_solver> {
  template <typename Model>
  static void apply(Model* model, TimeStep<implicit_solver>& tstep) {

    typename Model::Mesh* mesh        = model->getMesh();
    DofHandler*           dof_handler = model->getDofHandler();

    const _real stime = tstep.getStartTime();
    const _real etime = tstep.getEndTime();
    GEOXFEM_LOG_SEV(info) << "Solving time Step (" << stime << "--" << etime
                          << ")...";

    /* 1. update load information acoording to boundary condition */
    for(auto vertex : shapesIn(*mesh, 0))
      update_nodal_boundary_condition<CouplePolicy>::apply(tstep,
                                                           vertex->toVertex());
    size_t iter = 0;

    /* 2. entering iteration loop */
    do {
      model->setIter(iter++);
      /* 2.1 update dofs */
      dof_handler->update(*mesh);
      GEOXFEM_LOG_SEV(info)
        << "Iteration #" << iter << ", with NEQ " << dof_handler->getNEQ();
      /* 2.2 assemble kay and right hand side at current iteration step */
      assemble_equation<MethodPolicy, CouplePolicy>::apply(*mesh, *dof_handler,
                                                           tstep);
      /* 2.3 solve equation at current iteration step */
      tstep.solve();
      /* 2.4 update at each iteration */
      update_iteration<CouplePolicy, LoadPolicy>::apply(*mesh, *dof_handler, tstep);
      /* 2.5 output result */
      model->notify();
    } while(!tstep.ifConverged());

    /* 3. update at each time step */
    update_time_step<CouplePolicy, LoadPolicy>::apply(*mesh);

    return;
  }
}; /*}}}*/

}  // namespace model

namespace traits {

template <>
_real get_delta_time::apply</*{{{*/
                            model::TimeStep<model::implicit_solver>>(
  const model::TimeStep<model::implicit_solver>& tstep) {
  return tstep.getDurationTime();
} /*}}}*/

}  // namespace traits
}  // namespace geoxfem

#endif /* IMPLICIT_SOLVER_H */
